library(groundhog)

packages <- c(
  "lmtest",
  "sandwich",
  "tidyr",
  "xtable",
  "tidyverse",
  "broom",
  "ggplot2",
  "stargazer",
  "doBy",
  "tidyr",
  "foreign"
)
groundhog.library(packages, "2024-03-01")

## define function to rescale variables to range between 0 and 1
recode_0_1 <- function(x) {
  return((x - min(x, na.rm = T)) / (max(x, na.rm = T) - min(x, na.rm = T)))
}



## define function to make transparent colors in plots
# function from: https://www.dataanalytics.org.uk/make-transparent-colors-in-r/
t_col <- function(color, percent = 50, name = NULL) {
  #      color = color name
  #    percent = % transparency
  #       name = an optional name for the color
  
  ## Get RGB values for named color
  rgb.val <- col2rgb(color)
  
  ## Make new color using input color as base and alpha set by transparency
  t.col <- rgb(rgb.val[1], rgb.val[2], rgb.val[3],
               max = 255,
               alpha = (100 - percent) * 255 / 100,
               names = name
  )
  
  ## Save the color
  invisible(t.col)
}


## define function to estimate means by treatment condition and make plot parameters
dose_means <- function(data, dv, color_women, color_men, shape_women, shape_men) {
  dose_test <- c("control", "20", "30", "40", "50", "60", "70", "80")
  dv <- dv
  res2 <- as.data.frame(matrix(nrow = 15, ncol = 3))
  colnames(res2) <- c("dose_test", "coef", "se")
  res2
  gender <- c("men", "women")
  res2$gender <- c("control", rep(gender, 7))
  res2$dose_test <- c(
    "control", "20", "80", "30", "70", "40", "60", "50",
    "50", "60", "40", "70", "30", "80", "20"
  )
  res2
  
  if (dv == "ability") {
    for (i in 1:length(gender)) {
      for (j in 1:length(dose_test)) {
        if (dose_test[j] == "control") {
          m <- lm(ability ~ 1, data = data[data$dose == dose_test[j], ])
          temp <- coeftest(m, vcov = vcovCL, cluster = ~ID)
          res2[res2$dose_test == dose_test[j] & res2$gender == "control", "coef"] <- temp[1, "Estimate"]
          res2[res2$dose_test == dose_test[j] & res2$gender == "control", "se"] <- temp[1, "Std. Error"]
        }
        
        if (dose_test[j] != "control") {
          m <- lm(ability ~ 1, data = data[data$dose == dose_test[j] & data$gender_treatment == gender[i], ])
          temp <- coeftest(m, vcov = vcovCL, cluster = ~ID)
          res2[res2$dose_test == dose_test[j] & res2$gender == gender[i], "coef"] <- temp[1, "Estimate"]
          res2[res2$dose_test == dose_test[j] & res2$gender == gender[i], "se"] <- temp[1, "Std. Error"]
        }
      }
    }
  }
  
  if (dv == "rep_women") {
    for (i in 1:length(gender)) {
      for (j in 1:length(dose_test)) {
        if (dose_test[j] == "control") {
          m <- lm(rep_women ~ 1, data = data[data$dose == dose_test[j], ])
          temp <- coeftest(m, vcov = vcovCL, cluster = ~ID)
          res2[res2$dose_test == dose_test[j] & res2$gender == "control", "coef"] <- temp[1, "Estimate"]
          res2[res2$dose_test == dose_test[j] & res2$gender == "control", "se"] <- temp[1, "Std. Error"]
        }
        
        if (dose_test[j] != "control") {
          m <- lm(rep_women ~ 1, data = data[data$dose == dose_test[j] & data$gender_treatment == gender[i], ])
          temp <- coeftest(m, vcov = vcovCL, cluster = ~ID)
          res2[res2$dose_test == dose_test[j] & res2$gender == gender[i], "coef"] <- temp[1, "Estimate"]
          res2[res2$dose_test == dose_test[j] & res2$gender == gender[i], "se"] <- temp[1, "Std. Error"]
        }
      }
    }
  }
  
  res2$gender[is.na(res2$gender)] <- "control"
  res2$lb <- res2$coef - 1.96 * res2$se
  res2$ub <- res2$coef + 1.96 * res2$se
  res2$pch <- NA
  res2$pch[res2$gender == "women"] <- shape_women
  res2$pch[res2$gender == "men"] <- shape_men
  res2$color <- NA
  res2$color[res2$gender == "control"] <- "#E69F00"
  res2$color[res2$gender == "women"] <- color_women
  res2$color[res2$gender == "men"] <- color_men
  return(res2)
}

## ability
dose_means_prolific <- function(data, dv, color_women, color_men, shape_women, shape_men) {
  dose_test <- c("20", "30", "40", "50", "60", "70", "80")
  dv <- dv
  res2 <- as.data.frame(matrix(nrow = 14, ncol = 3))
  colnames(res2) <- c("dose_test", "coef", "se")
  res2
  gender <- c("men", "women")
  res2$gender <- c(rep(gender, 7))
  res2$dose_test <- c(
    "20", "80", "30", "70", "40", "60", "50",
    "50", "60", "40", "70", "30", "80", "20"
  )
  res2
  
  if (dv == "e1_agency_confidence") {
    for (i in 1:length(gender)) {
      for (j in 1:length(dose_test)) {
        m <- lm(e1_agency_confidence ~ 1, data = data[data$pct_dose == dose_test[j] & data$exp_1_gender == gender[i], ])
        temp <- coeftest(m, vcov = vcovHC(m, type = "HC1"))
        res2[res2$dose_test == dose_test[j] & res2$gender == gender[i], "coef"] <- temp[1, "Estimate"]
        res2[res2$dose_test == dose_test[j] & res2$gender == gender[i], "se"] <- temp[1, "Std. Error"]
      }
    }
  }
  
  if (dv == "e1_rep_women") {
    for (i in 1:length(gender)) {
      for (j in 1:length(dose_test)) {
        if (dose_test[j] != "control") {
          m <- lm(e1_rep_women ~ 1, data = data[data$pct_dose == dose_test[j] & data$exp_1_gender == gender[i], ])
          temp <- coeftest(m, vcov = vcovHC(m, type = "HC1"))
          res2[res2$dose_test == dose_test[j] & res2$gender == gender[i], "coef"] <- temp[1, "Estimate"]
          res2[res2$dose_test == dose_test[j] & res2$gender == gender[i], "se"] <- temp[1, "Std. Error"]
        }
      }
    }
  }
  
  res2$gender[is.na(res2$gender)] <- "control"
  res2$lb <- res2$coef - 1.96 * res2$se
  res2$ub <- res2$coef + 1.96 * res2$se
  res2$pch <- NA
  res2$pch[res2$gender == "women"] <- shape_women
  res2$pch[res2$gender == "men"] <- shape_men
  res2$color <- NA
  res2$color[res2$gender == "control"] <- "#E69F00"
  res2$color[res2$gender == "women"] <- color_women
  res2$color[res2$gender == "men"] <- color_men
  return(res2)
}